#include "typechecker.hpp"
#include "symtable.hpp"
#include "ast.hpp"
#include "generic_cleaner.hpp"

#include <cassert>
#include <iostream>
using namespace std;

bool TypeChecker::has_error() {
    return _has_error;
}

TypeTable* TypeChecker::type_table() {
    return _sym_table->type_table();
}

Node* TypeChecker::visit(Node* n) {
    return n->accept(this);
}

Node* TypeChecker::visit(PrintNode* n) {
    visit(n->value());
    return n;
}

Node* TypeChecker::visit(ReturnNode* n) {
    visit(n->value());
    if (!_sym_table->is_func()) {
        cout << "Can not use return outside of function" << endl;
        _has_error = true;
        return n;
    }

    Type* ret_type = _sym_table->ret_type();
    Type* value_type = n->value()->get_type(type_table());

    if (!unify(ret_type, value_type)) {
        cerr << "return type mismatch: " << ret_type->to_string()
            << " and " << value_type->to_string() << endl;
        _has_error = true;
    }

    n->set_type(ret_type);
    return n;
}

Node* TypeChecker::visit(BinaryOp* op) {
    visit(op->left());
    Type* left_type = op->left()->get_type(type_table());

    visit(op->right());
    Type* right_type = op->right()->get_type(type_table());

    switch(op->op_type()) {
    case AST_OP_ADD:
        if (!left_type->support_operator(OPERATOR_ADD, right_type)) {
            cerr << "Unsupported operator '+' for type " << left_type->to_string()
                << " and " << right_type->to_string() << endl;
            _has_error = true;
            return op;
        }

        op->set_type(left_type);
        break;
    case AST_OP_SUB:
        if (!left_type->support_operator(OPERATOR_SUB, right_type)) {
            cerr << "Unsupported operator '-' for type " << left_type->to_string()
                << " and " << right_type->to_string() << endl;
            _has_error = true;
            return op;
        }

        op->set_type(left_type);
        break;
    case AST_OP_MUL:
        if (!left_type->support_operator(OPERATOR_MUL, right_type)) {
            cerr << "Unsupported operator '*' for type " << left_type->to_string()
                << " and " << right_type->to_string() << endl;
            _has_error = true;
            return op;
        }

        op->set_type(left_type);
        break;
    case AST_OP_DIV:
        if (!left_type->support_operator(OPERATOR_DIV, right_type)) {
            cerr << "Unsupported operator '/' for type " << left_type->to_string()
                << " and " << right_type->to_string() << endl;
            _has_error = true;
            return op;
        }

        op->set_type(left_type);
        break;
    case AST_OP_MOD:
        // TODO: check left and right are int.
        // if left is string, then right must be tuple.
        break;
    case AST_OP_BIT_AND:
        break;
    case AST_OP_BIT_OR:
        break;
    case AST_OP_BIT_XOR:
        break;
    case AST_OP_LOG_OR:
        break;
    case AST_OP_LOG_AND:
        break;
    case AST_OP_LEFT_SHIFT:
        break;
    case AST_OP_RIGHT_SHIFT:
        break;
    case AST_OP_COMPOSE: {
        if (!left_type->is_arrow() || !right_type->is_arrow()) {
            cerr << "compose type must be function, but get " << left_type->to_string()
                << " and " << right_type->to_string() << endl;
            _has_error = true;
            return op;
        }
        
        ArrowType* left_at = dynamic_cast<ArrowType*>(left_type);
        ArrowType* right_at = dynamic_cast<ArrowType*>(right_type);
        Type* left_arg_type = left_at->src();
        Type* left_ret_type = left_at->dst();

        if (!unify(left_ret_type, right_at->src())) {
            cerr << "compose type mismatch: " << left_arg_type->to_string()
                << " and " << right_at->src()->to_string() << endl;
            _has_error = true;
            return op;
        }

        op->set_type(new ArrowType(left_arg_type, right_at->dst()));

        break;
    }
    default:
        printf("Unknown binary op %d\n", op->op_type());
        return op;
    }
    return op;
}

Node* TypeChecker::visit(ListNode* n) {
    for (auto it = n->node_list()->begin(); 
        it != n->node_list()->end(); it++) {
        visit(*it);
    }
    return n;
}

Node* TypeChecker::visit(StmtsNode* n) {
    for (auto it = n->node_list()->begin(); 
        it != n->node_list()->end(); it++) {
        visit(*it);
    }
    return n;
}

Node* TypeChecker::visit(ConstInt* n) {
    return n;
}

Node* TypeChecker::visit(ConstDouble* n) {
    return n;
}

Node* TypeChecker::visit(BuiltinConst* n) {
    return n;
}

Node* TypeChecker::visit(ConstString* n) {
    return n;
}

Node* TypeChecker::visit(ConstArray* n) {
    if (n->arr()->size() == 0) {
        return n;
    }

    visit(n->arr()->get(0));
    Type* elem_type = n->arr()->get(0)->get_type(type_table());
    for (int i = 1; i < n->arr()->size(); i++) {
        visit(n->arr()->get(i));
        Type* t = n->arr()->get(i)->get_type(type_table());
        if (!unify(elem_type, t)) {
            cerr << "List elements must be compatible." << endl;
            _has_error = true;
            return nullptr;
        }
    }

    n->set_type(new ListType(elem_type));
    return n;
}

Node* TypeChecker::visit(ImportNode* n) {
    return n;
}

Node* TypeChecker::visit(ImplementNode* n) {
    return n;
}

Node* TypeChecker::visit(IfNode* n) {
    visit(n->cond());
    visit(n->then_part());

    if (n->else_part() != NULL) {
        visit(n->else_part());
    }
    return n;
}

Node* TypeChecker::visit(WhileNode* n) {
    visit(n->cond());
    visit(n->body());
    return n;
}

Node* TypeChecker::visit(CmpNode* n) {
    visit(n->left());
    visit(n->right());
    if (n->cmp_op() == CmpNode::CMP_LT) {
    }

    return n;
}

Node* TypeChecker::visit(LogicNotNode* n) {
    visit(n->value());
    return n;
}

Node* TypeChecker::visit(AssignNode* n) {
    visit(n->left());
    Type* left_type = n->left()->get_type(type_table());

    if (!left_type) {
        cerr << "undeclared variable in assignment." << endl;
        _has_error = true;
        return n;
    }

    visit(n->right());
    Type* right_type = n->right()->get_type(type_table());

    if (!right_type) {
        cerr << "undeclared variable in assignment." << endl;
        _has_error = true;
        return n;
    }

    if (!unify(left_type, right_type)) {
        cerr << "Incompatible assigning from " << right_type->to_string()
            << " to " << left_type->to_string() << endl;
        _has_error = true;
    }

    return n;
}

Node* TypeChecker::visit(VarNode* n) {
    int index = _sym_table->get_name_index(n->name());
    Type* var_type = nullptr;
    if (index >= 0) {
        var_type = _sym_table->get_name_type(n->name());
        n->set_type(var_type);
        return n;
    }

    index = _sym_table->get_var_name_index(n->name());
    if (index >= 0) {
        var_type = _sym_table->get_var_name_type(n->name());
        n->set_type(var_type);
        return n;
    }

    index = _sym_table->get_local_index(n->name());
    if (index >= 0) {
        var_type = _sym_table->get_local_type(n->name());
        n->set_type(var_type);
        return n;
    }

    index = _sym_table->get_cell_var_index(n->name());
    if (index >= 0) {
        var_type = _sym_table->get_cell_var_type(n->name());
        n->set_type(var_type);
        return n;
    }

    index = _sym_table->get_free_var_index(n->name());
    if (index >= 0) {
        var_type = _sym_table->get_free_var_type(n->name());
        n->set_type(var_type);
        return n;
    }

    // find name along the sym table chain.
    SymbolTable* pst = _sym_table->parent();
    while (pst != nullptr) {
        index = pst->get_name_index(n->name());
        if (index >= 0) {
            var_type = pst->get_name_type(n->name());
            _sym_table->add_name(n->name(), var_type);
            break;
        }

        if (!pst->is_func()) {
            pst = pst->parent();
            continue;
        }

        // If pst is func, try to find variable in varnames.
        index = pst->get_var_name_index(n->name());
        if (index >= 0) {
            var_type = pst->get_var_name_type(n->name());

            /* For those functions that variable is not defined in,
             * this variable is free, else it is cell.
             */
            SymbolTable* begin = _sym_table;
            while (begin != pst) {
                if (begin->get_free_var_index(n->name()) < 0) {
                    begin->add_free_var(n->name(), var_type);
                }
                begin = begin->parent();
            }

            if (pst->get_cell_var_index(n->name()) < 0) {
                pst->add_cell_var(n->name());
            }
            break;
        }
        pst = pst->parent();
    }

    if (index == -1) {
        var_type = type_table()->get_expanded_type(n->name());
        if (!var_type) {
            cerr << "undeclared variable '" << n->name() << "'." << endl;
            _has_error = true;
        }
    }

    n->set_type(var_type);

    return n;
}

Node* TypeChecker::visit(DotNode* n) {
    visit(n->name());

    Type* name_type = n->name()->get_type(type_table());
    if (!name_type) {
        cerr << "Unknown type for variable '" << n->name()->to_string() << "'." << endl;
        _has_error = true;
        return nullptr;
    }

    if (name_type->is_record()) {
        RecordType* rt = dynamic_cast<RecordType*>(name_type);
        VarNode* field_node = dynamic_cast<VarNode*>(n->field());
        assert(field_node);

        if (rt->fields().find(field_node->name()) != rt->fields().end()) {
            n->set_type(rt->get_field_type(field_node->name()));
        }
        else if (rt->get_method_type(field_node->name())) {
            n->set_type(rt->get_method_type(field_node->name()));
        }
        else {
            cerr << "Unknown field '" << field_node->name() << "' in record type '" << name_type->to_string() << "'." << endl;
            _has_error = true;
        }
    }
    else {
        cerr << "Cant not getattr on '" << n->name()->to_string() << "'." << endl;
        _has_error = true;
        return nullptr;
    }
    
    return n;
}

Node* TypeChecker::visit(SequenceNode* n) {
    visit(n->left());

    if (!unify(n->left()->get_type(type_table()), IntType::get_instance())) {
        cerr << "Sequence must start as an integer." << endl;
        _has_error = true;
        return nullptr;
    }

    visit(n->right());
    if (!unify(n->right()->get_type(type_table()), IntType::get_instance())) {
        cerr << "Sequence must end as an integer." << endl;
        _has_error = true;
        return nullptr;
    }

    n->set_type(IntType::get_instance());
    return n;
}

Node* TypeChecker::visit(CallNode* n) {
    Node* callee = visit(n->callee());

    Type* callee_type = callee->get_type(type_table());
    if (!callee_type) {
        cerr << "undeclared function." << endl;
        _has_error = true;
        return n;
    }

    if (callee_type->is_list()) {
        return n;
    }

    if (callee_type->is_record()) {
        RecordType* rt = dynamic_cast<RecordType*>(callee_type);
        if (rt->fields().size() != n->args()->size()) {
            cerr << "record type has " << rt->fields().size() << " fields, but " << n->args()->size() << " arguments are provided." << endl;
            _has_error = true;
            return n;
        }

        int i = 0;
        for (auto it = n->args()->node_list()->begin();
            it != n->args()->node_list()->end(); it++) {
            visit(*it);
            if (!unify((*it)->get_type(type_table()), rt->field_types()[i])) {
                cerr << "record type mismatch: " << (*it)->get_type(type_table())->to_string() 
                    << " and " << rt->field_types()[i]->to_string() << endl;
                _has_error = true;
                break;
            }
            i++;
        }

        n->set_type(rt);

        return n;
    }

    callee_type = apply_args(callee_type, n->args())->substitute(type_table());
    n->set_type(callee_type);

    return n;
}

Type* TypeChecker::apply_args(Type* callee_type, ListNode* args) {
    ArrowType* at = dynamic_cast<ArrowType*>(callee_type);
    if (!at) {
        cerr << "Callee is not a function." << endl;
        _has_error = true;
        return nullptr;
    }

    if (args->size() == 0) {
        return at->try_apply(UnitType::get_instance());
    }

    // apply all arguments one by one.
    for (auto it = args->node_list()->begin();
        it != args->node_list()->end(); it++) {
        visit(*it);

        at = dynamic_cast<ArrowType*>(callee_type);
        if (!at) {
            cerr << "Callee is not a function." << endl;
            _has_error = true;
            return nullptr;
        }

        Type* arg_type = (*it)->get_type(type_table());
        if (!unify(at->src(), arg_type)) {
            cerr << "Argument type mismatch: " << at->src()->to_string() 
                << " and " << arg_type->to_string() << endl;
            _has_error = true;
            break;
        }

        callee_type = at->try_apply(arg_type);
        if (!callee_type) {
            cout << "Call function with wrong arguments." << endl;
            _has_error = true;
            break;
        }
    }

    return callee_type;
}

Node* TypeChecker::visit(TupleNode* n) {
    ListNode* args = n->args();
    for (auto it = args->node_list()->begin();
        it != args->node_list()->end(); it++) {
        visit(*it);
    }
    return n;
}

Node* TypeChecker::visit(TypeTupleNode* n) {
    return n;
}

Node* TypeChecker::visit(FuncDefNode* n) {
    if (n->has_generic_args()) {
        return n;
    }

    _sym_table = n->sym_table();
    visit(n->args());
    if (n->body()) {
        visit(n->body());
    }

    _sym_table = _sym_table->parent();

    return n;
}

Node* TypeChecker::visit(RecordDefNode* n) {
    return n;
}

Node* TypeChecker::visit(TraitDefNode* n) {
    for (auto it = n->methods()->node_list()->begin();
        it != n->methods()->node_list()->end(); it++) {
        visit(*it);
    }
    return n;
}

Type* TypeChecker::instantiate_generic_func(FuncDefNode* n, ListNode* type_args) {
    /* 清理上一次实例化时的状态 */
    if (!n->has_generic_args()) {
        cerr << "Generic function has no generic arguments." << endl;
        _has_error = true;
        return nullptr;
    }

    // point to use point.
    n->sym_table()->set_parent(_sym_table);
    _sym_table = n->sym_table();
     
    {
        GenericCleaner gc;
        gc.visit(n);
    }

    // Process type arguments - these are already parsed as type nodes
    int i = 0;
    for (auto it = type_args->begin(); it != type_args->end(); it++) {
        visit(*it); // Visit the type node
        type_table()->set_generic_name(i++, (*it)->get_type(type_table()));
    }

    visit(n->args());
    visit(n->body());

    Type* func_type = build_func_type(n);
    func_type = func_type->substitute(type_table());

    /* Restore to parent's scope */
    _sym_table = _sym_table->parent();

    return func_type;
}

Node* TypeChecker::visit(SubscrNode* n) {
    visit(n->name());

    Type* var_type = n->name()->get_type(type_table());
    if (!var_type) {
        cerr << "undeclared variable." << endl;
        _has_error = true;
        // TODO : this is a workaround for record.
        return n->name();
    }
    
    if (var_type->is_poly()) {
        // Handle generic instantiation
        // The parser has already parsed this as type arguments
        ListNode* type_args = dynamic_cast<ListNode*>(n->index());
        if (!type_args) {
            cerr << "Generic instantiation expects type arguments list" << endl;
            _has_error = true;
            return nullptr;
        }

        VarNode* name_node = dynamic_cast<VarNode*>(n->name());
        if (!name_node) {
            cerr << "Generic function name must be a simple identifier" << endl;
            _has_error = true;
            return nullptr;
        }

        FuncDefNode* func_def = dynamic_cast<FuncDefNode*>(_sym_table->get_subtree_by_name(name_node->name()));
        if (!func_def) {
            cerr << "Cannot find generic function definition for " << name_node->name() << endl;
            _has_error = true;
            return nullptr;
        }

        Type* t = instantiate_generic_func(func_def, type_args);
        n->set_type(t);
        return n;
    }

    if (var_type->is_type_func()) {
        // nothing to do. Check arguments in CallNode.
        visit(n->index());
        TypeFunction* tf = dynamic_cast<TypeFunction*>(var_type);
        ListNode* type_args = dynamic_cast<ListNode*>(n->index());
        if (!type_args) {
            cerr << "TypedVisitor : generic instantiation expects type arguments list" << endl;
            _has_error = true;
            return nullptr;
        }
    
        // 一个record实例化的本质就是函数调用：入参是类型，输出是一个新的类型，
        // 所以这里需要一个环境来记录类型参数的替换，就像函数调用的栈空间一样，
        // 调用结束以后，栈空间被销毁，类型参数的替换也被销毁。
        TypeTable env(type_table());
    
        for (int i = 0; i < type_args->size(); i++) {
            // must be type var.
            TypeVar* tv = dynamic_cast<TypeVar*>(tf->args()[i]);
            env.add_substitution(tv->name(), type_args->get(i)->get_type(type_table()));
        }
        
        env.add_substitution("Self", var_type);
        Type* new_type = tf->expand(&env);
    
        n->set_type(new_type);
        return n;
    }

    // Handle array access or other subscript operations
    // The parser has already parsed this as an expression
    visit(n->index());
    Type* index_type = n->index()->get_type(type_table());
    
    // Re-get the left operand type
    visit(n->name());
    Type* left_type = n->name()->get_type(type_table());

    left_type = left_type->substitute(type_table());
    
    if (left_type->is_list() || left_type->is_tuple()) {
        if (index_type != IntType::get_instance()) {
            cerr << "Array index must be integer type" << endl;
            _has_error = true;
            return nullptr;
        }
        if (left_type->is_list()) {
            ListType* list_type = dynamic_cast<ListType*>(left_type);
            n->set_type(list_type->elm_type());
        } else if (left_type->is_tuple()) {
            ConstInt* index_node = dynamic_cast<ConstInt*>(n->index());
            if (!index_node) {
                cerr << "Tuple index must be a const integer." << endl;
                _has_error = true;
                return nullptr;
            }
            TupleType* tuple_type = dynamic_cast<TupleType*>(left_type);
            n->set_type(tuple_type->types()[index_node->value()]);
        }

        return n;
    } else {
        cerr << "Subscript operation on non-array type: " << left_type->to_string() << endl;
        _has_error = true;
    }

    return nullptr;
}

Type* TypeChecker::get_param_type(ParamDefNode* n) {
    return n->get_type(type_table());
}

Type* TypeChecker::build_func_type(FuncDefNode* n) {
    ListNode* args = n->args();

    Type* func_type = n->ret_type();
    if (!func_type) { 
        cerr << "TypeChecker : Function '" << n->name() << "' has no return type." << endl;
        _has_error = true;
        return nullptr;
    }

    if (args->size() == 0) {
        return new ArrowType(UnitType::get_instance(), func_type);
    }

    for (auto it = args->node_list()->rbegin();
        it != args->node_list()->rend(); it++) {
        Type* t = get_param_type(dynamic_cast<ParamDefNode*>(*it));
        func_type = new ArrowType(t, func_type);
    }

    return func_type;
}

Node* TypeChecker::visit(ParamDefNode* n) {
    return n;
}

Node* TypeChecker::visit(VarDefNode* n) {
    Type* declare_type = n->get_type(type_table());
    Type* value_type = nullptr;

    if (n->init_value()) {
        visit(n->init_value());
        value_type = n->init_value()->get_type(type_table());
        if (!value_type) {
            cerr << "undeclared type for variable '" << n->name() << "'." << endl;
            _has_error = true;
        }
    }

    // check type rules.
    if (value_type) {
        if (!unify(declare_type, value_type)) {
            cerr << "types conflict for " << n->name()  << endl;
            cerr << "declare type: " << declare_type->to_string() << endl;
            cerr << "value type: " << value_type->to_string() << endl;
            _has_error = true;
        }
    }

    return n;
}

Node* TypeChecker::visit(TypeVarNode* n) {
    return nullptr;
}

Node* TypeChecker::visit(TypeVarDefNode* n) {
    return nullptr;
}

Node* TypeChecker::visit(TypeFuncDefNode* n) {
    TypeFunction* tf = new TypeFunction(nullptr);

    enter_type_scope();
    for (auto it = n->args()->node_list()->begin(); 
        it != n->args()->node_list()->end(); it++) {
        TypeVarNode* tvn = dynamic_cast<TypeVarNode*>(*it);
        visit(tvn);
        tf->args().push_back(tvn->get_type(type_table()));
    }

    if (n->body()) {
        visit(n->body());
        tf->set_body(n->body()->get_type(type_table()));
    }

    exit_type_scope();

    n->set_type(tf);
    type_table()->add(n->name(), tf);

    return nullptr;
}

void TypeChecker::enter_type_scope() {
    _sym_table = new SymbolTable(_sym_table, true);
}

void TypeChecker::exit_type_scope() {
    _sym_table = _sym_table->parent();
}

Node* TypeChecker::visit(TypeCallNode* n) {
    return n;
}

Node* TypeChecker::visit(TypeConstInt* n) {
    return nullptr;
}

Node* TypeChecker::visit(ArrowTypeNode* n) {
    return nullptr;
}

Node* TypeChecker::visit(LambdaDef* n) {
    /* In function scope. */
    _sym_table = n->sym_table();
    _sym_table->set_arg_count(1);

    visit(n->param());
    visit(n->body());

    Type* func_type = new ArrowType(n->param()->get_type(type_table()), 
        _sym_table->ret_type());
    n->set_type(func_type);
    _sym_table = _sym_table->parent();

    return n;
}

/* pseudo nodes */
Node* TypeChecker::visit(ArrayTypeNode* n) {
    return nullptr;
}

Node* Evaluator::visit(Node* n) {
    return n->accept(this);
}

Node* Evaluator::visit(TypeConstInt* n) {
    _value.cvt = CVT_INT;
    _value.value.int_value = n->value;

    return nullptr;
}

bool TypeChecker::unify(Type* left, Type* right) {
    if (left == right) {
        return true;
    }
    
    if (left->equals(right)) {
        return true;
    }

    if (left->is_list() && right->is_list()) {
        ListType* left_lt = dynamic_cast<ListType*>(left);
        ListType* right_lt = dynamic_cast<ListType*>(right);
        return unify(left_lt->elm_type(), right_lt->elm_type());
    }

    if (left->is_tuple() && right->is_tuple()) {
        TupleType* left_tt = dynamic_cast<TupleType*>(left);
        TupleType* right_tt = dynamic_cast<TupleType*>(right);
        if (left_tt->types().size() != right_tt->types().size()) {
            return false;
        }

        for (int i = 0; i < left_tt->types().size(); i++) {
            if (!unify(left_tt->types()[i], right_tt->types()[i])) {
                return false;
            }
        }

        return true;
    }

    if (left->is_arrow() && right->is_arrow()) {
        ArrowType* left_at = dynamic_cast<ArrowType*>(left);
        ArrowType* right_at = dynamic_cast<ArrowType*>(right);
        return unify(left_at->src(), right_at->src()) && unify(left_at->dst(), right_at->dst());
    }

    if (left->is_record() && right->is_record()) {
        RecordType* left_rt = dynamic_cast<RecordType*>(left);
        RecordType* right_rt = dynamic_cast<RecordType*>(right);
        if (left_rt->name() != right_rt->name()) {
            return false;
        }

        assert(left_rt->fields().size() == right_rt->fields().size());
        assert(left_rt->generic_args().size() == right_rt->generic_args().size());

        for (auto it = left_rt->fields().begin(); it != left_rt->fields().end(); it++) {
            if (!unify(it->second, right_rt->fields()[it->first])) {
                return false;
            }
        }

        for (auto it = left_rt->generic_args().begin(); it != left_rt->generic_args().end(); it++) {
            if (!unify(it->second, right_rt->get_generic_arg(it->first))) {
                return false;
            }
        }

        return true;
    }

    Type* left_parent = left->is_type_var() ? 
            dynamic_cast<TypeVar*>(left)->env()->ufset_find(left) : left;
    Type* right_parent = right->is_type_var() ? 
            dynamic_cast<TypeVar*>(right)->env()->ufset_find(right) : right;

    if (left_parent->equals(right_parent)) {
        return true;
    }

    if (left_parent->is_arrow() && right_parent->is_arrow()) {
        return unify(left_parent, right_parent);
    }

    if (left_parent->is_record() && right_parent->is_record()) {
        return unify(left_parent, right_parent);
    }

    if (left_parent->is_type_var()) {
        TypeVar* left_tv = dynamic_cast<TypeVar*>(left_parent);
        left_tv->env()->ufset_union(left_tv, right_parent);
    }
    else if (right_parent->is_type_var()) {
        TypeVar* right_tv = dynamic_cast<TypeVar*>(right_parent);
        right_tv->env()->ufset_union(right_tv, left_parent);
    }
    else {
        cerr << "TypeTable::ufset_union : type is not type var." << endl;
        cerr << "t1: " << left_parent->to_string() << endl;
        cerr << "t2: " << right_parent->to_string() << endl;
        assert(false);
    }

    return true;
}